import os

import dotenv
import requests
from langchain_community.utilities import GoogleSerperAPIWrapper

from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.tools import tool, Tool, BaseTool
from internal.lib.helper import add_attribute

dotenv.load_dotenv()


class GoogleSerperArgsSchema(BaseModel):
    """谷歌serper搜索参数描述"""
    query: str = Field(description="需要检索查询的语句")


@tool
def bocha_web_search_tool(query: str, count: int = 8) -> str:
    """
    使用Bocha Web Search API进行联网搜索，返回搜索结果的字符串。

    参数:
    - query: 搜索关键词
    - count: 返回的搜索结果数量

    返回:
    - 搜索结果的字符串形式
    """
    url = 'https://api.bochaai.com/v1/web-search'
    headers = {
        'Authorization': f'Bearer {os.getenv("BOCHA_API_KEY")}',  # 请替换为你的API密钥
        'Content-Type': 'application/json'
    }
    data = {
        "query": query,
        "freshness": "noLimit",  # 搜索的时间范围，例如 "oneDay", "oneWeek", "oneMonth", "oneYear", "noLimit"
        "summary": True,  # 是否返回长文本摘要总结
        "count": count
    }

    response = requests.post(url, headers=headers, json=data)

    if response.status_code == 200:
        # 返回给大模型的格式化的搜索结果文本
        # 可以自己对博查的搜索结果进行自定义处理
        return str(response.json())
    else:
        raise Exception(f"API请求失败，状态码: {response.status_code}, 错误信息: {response.text}")

@add_attribute("args_schema", GoogleSerperArgsSchema)
def bocha_search(**kwargs) -> BaseTool:
    """谷歌Serper搜索工具"""
    return Tool(
        name="BochaWebSearch",
        func=bocha_web_search_tool,
        description="使用Bocha Web Search API进行网络搜索",
        args_schema=GoogleSerperArgsSchema,
        api_wrapper=GoogleSerperAPIWrapper(),
    )